import React, { useState } from "react";
import { Transfer, Tree } from "antd";
import type { TransferDirection, TransferItem } from "antd/es/transfer";
import type { DataNode, TreeProps } from "antd/es/tree";

interface TreeTransferProps {
  dataSource: DataNode[];
  targetKeys: React.Key[];
  onChange: (
    targetKeys: React.Key[],
    direction: TransferDirection,
    moveKeys: React.Key[]
  ) => void;
}

// 判断节点是否被选中
const isChecked = (selectedKeys: React.Key[], eventKey: React.Key) =>
  selectedKeys.includes(eventKey);

// 生成树结构
const generateTree = (
  treeNodes: DataNode[] = [],
  checkedKeys: React.Key[] = []
): DataNode[] =>
  treeNodes.map(({ children, ...props }) => ({
    ...props,
    disabled: checkedKeys.includes(props.key as string),
    children: generateTree(children, checkedKeys),
  }));

const TreeTransfer = ({
  dataSource,
  targetKeys,
  onChange,
}: TreeTransferProps) => {
  const transferDataSource: TransferItem[] = [];

  // 扁平化树结构
  function flatten(list: DataNode[] = []) {
    list.forEach((item) => {
      transferDataSource.push(item as TransferItem);
      flatten(item.children);
    });
  }
  flatten(dataSource);

  // 处理双击事件 - 将节点穿梭到右侧
  const handleDoubleClick: TreeProps["onDoubleClick"] = (event, node) => {
    const { key } = node;
    // 检查节点是否已在右侧，不在则添加
    if (!targetKeys.includes(key)) {
      const newTargetKeys = [...targetKeys, key];
      onChange(newTargetKeys, "right", [key]);
    }
  };

  return (
    <Transfer
      onChange={(targetKeys, direction, moveKeys) => {
        onChange(targetKeys, direction, moveKeys);
      }}
      targetKeys={targetKeys}
      dataSource={transferDataSource}
      className="tree-transfer"
      render={(item) => item.title!}
      showSelectAll={false}
    >
      {({ direction, onItemSelect, selectedKeys }) => {
        if (direction === "left") {
          const checkedKeys: React.Key[] = [...selectedKeys, ...targetKeys];
          return (
            <Tree
              blockNode
              checkable
              checkStrictly
              defaultExpandAll
              checkedKeys={checkedKeys}
              treeData={generateTree(dataSource, targetKeys)}
              onCheck={(_, { node: { key } }) => {
                onItemSelect(key as string, !isChecked(checkedKeys, key));
              }}
              onSelect={(_, { node: { key } }) => {
                onItemSelect(key as string, !isChecked(checkedKeys, key));
              }}
              // 添加双击事件处理
              onDoubleClick={handleDoubleClick}
            />
          );
        }
      }}
    </Transfer>
  );
};

const treeData: DataNode[] = [
  { key: "0-0", title: "0-0" },
  {
    key: "0-1",
    title: "0-1",
    children: [
      { key: "0-1-0", title: "0-1-0" },
      { key: "0-1-1", title: "0-1-1" },
    ],
  },
  { key: "0-2", title: "0-2" },
];

const App: React.FC = () => {
  const [targetKeys, setTargetKeys] = useState<React.Key[]>([]);

  const onChange = (
    keys: React.Key[],
    direction: TransferDirection,
    moveKeys: React.Key[]
  ) => {
    setTargetKeys(keys);
  };

  return (
    <TreeTransfer
      dataSource={treeData}
      targetKeys={targetKeys}
      onChange={onChange}
    />
  );
};

export default App;
